import numpy as np
import scipy.sparse as sp
import torch
'''
one-hot编码
将所有由字符串表示的标签数组用set保存，set的重要特征是元素没有重复
随后为每个标签分配一个编号，创建一个单位矩阵，单位矩阵的每一行对应一个one-hot向量，也就是np.identity(len(classes))[i, :]
将每个数据对应的标签表示成的one-hot向量，类型为numpy数组
'''


def encode_onehot(labels):
    classes = set(labels)
    classes_dict = {c: np.identity(len(classes))[i, :] for i, c in
                    enumerate(classes)}
    labels_onehot = np.array(list(map(classes_dict.get, labels)),
                             dtype=np.int32)
    return labels_onehot


def load_data(path="data/cora/", dataset="cora"):
    """Load citation network dataset (cora only for now)"""
    print('Loading {} dataset...'.format(dataset))

    # 读取content文件，快速将txt文本内容转换成numpy数组
    # idx_features_labels[:, 0]、idx_features_labels[:, 1:-1]、idx_features_labels[:, -1]分别代表id features labels
    idx_features_labels = np.genfromtxt("{}{}.content".format(path, dataset),
                                        dtype=np.dtype(str))
    # 将第二部分features存储为稀疏矩阵
    features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    # 将第三部分labels表示成one-hot向量的形式
    labels = encode_onehot(idx_features_labels[:, -1])

    # 构建网络
    # 将所有节点的id表示成整数类型的numpy数组
    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    # content中节点不是按顺序排列的，建立哈希表，每一项为id:number，节点id对应的编号为number
    idx_map = {j: i for i, j in enumerate(idx)}
    # 读取cites文件，转换成numpy中数组，形状是（edge_num,2），每一行表示一条边两个端点的idx
    edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset),
                                    dtype=np.int32)
    # 将边的2个端点id换成编号number
    edges = np.array(list(map(idx_map.get, edges_unordered.flatten())),
                     dtype=np.int32).reshape(edges_unordered.shape)
    # 根据给出的下标、数据和形状构建矩阵，网络中有多少边矩阵中 就有多少个1
    # 先创建长度为edge_num的全1矩阵，每个1填充的位置是一条边中两个端点的编号(edges[:, 0], edges[:, 1])，矩阵形状(node_size, node_size)
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])),
                        shape=(labels.shape[0], labels.shape[0]),
                        dtype=np.float32)

    # 构建对称矩阵A^，后面构建model会用
    # 上面的adj是有向图的邻接矩阵，转成无向图的对称邻接矩阵
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
    # 分别对features特征矩阵和adj邻接矩阵做标准化
    features = normalize(features)
    adj = normalize(adj + sp.eye(adj.shape[0]))

    #分别构建训练集、验证集、测试集，并创建特征矩阵、标签向量和邻接矩阵的tensor，用来做模型的输入。
    idx_train = range(140)
    idx_val = range(200, 500)
    idx_test = range(500, 1500)

    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.LongTensor(np.where(labels)[1])
    adj = sparse_mx_to_torch_sparse_tensor(adj)

    idx_train = torch.LongTensor(idx_train)
    idx_val = torch.LongTensor(idx_val)
    idx_test = torch.LongTensor(idx_test)

    return adj, features, labels, idx_train, idx_val, idx_test


def normalize(mx):
    """
    Row-normalize sparse matrix
    首先对每一行求和得到rowsum；求倒数得到r_inv；如果某一行全为0，则r_inv算出来会等于无穷大，将这些行的r_inv置为0；
    构建对角元素为r_inv的对角矩阵；用对角矩阵与原始矩阵的点积起到标准化的作用，原始矩阵中每一行元素都会与对应的r_inv相乘。
    """
    rowsum = np.array(mx.sum(1))
    r_inv = np.power(rowsum, -1).flatten()
    r_inv[np.isinf(r_inv)] = 0.
    r_mat_inv = sp.diags(r_inv)
    mx = r_mat_inv.dot(mx)
    return mx


def accuracy(output, labels):
    preds = output.max(1)[1].type_as(labels)
    correct = preds.eq(labels).double()
    correct = correct.sum()
    return correct / len(labels)


def sparse_mx_to_torch_sparse_tensor(sparse_mx):
    """Convert a scipy sparse matrix to a torch sparse tensor."""
    sparse_mx = sparse_mx.tocoo().astype(np.float32)
    indices = torch.from_numpy(
        np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
    values = torch.from_numpy(sparse_mx.data)
    shape = torch.Size(sparse_mx.shape)
    return torch.sparse.FloatTensor(indices, values, shape)
